using System;
using System.Collections.Generic;
using System.IO;

using Multiplicity.Runtime.Serialization.Binary;

namespace Multiplicity.Reflection {
    public sealed class PEHeader {

        public PEHeaderType Type;

        public Version LinkerVersion;
        public Version OperatingSystemVersion;
        public Version ImageVersion;
        public Version SubsystemVersion;
        public uint SizeOfCode;
        public uint SizeOfInitializedData;
        public uint SizeOfUninitializedData;
        public uint AddressOfEntryPoint;
        public uint BaseOfCode;
        public uint BaseOfData;

        public ulong ImageBase;
        public uint SectionAlignment;
        public uint FileAlignment;
        public uint Win32VersionValue;
        public uint SizeOfImage;
        public uint SizeOfHeaders;
        public uint Checksum;
        public ushort Subsystem;
        public ushort DllCharacteristics;
        public ulong SizeOfStackReserve;
        public ulong SizeOfStackCommit;
        public ulong SizeOfHeapReserve;
        public ulong SizeOfHeapCommit;
        public uint LoaderFlags;

        public PEDataDirectory[] DataDirectories;
        
        public PEHeader() {
        }

        public static PEHeader Parse(BinaryParser br) {
            PEHeaderType signature = (PEHeaderType)br.ReadUInt16();
            if (!Enum.IsDefined(typeof(PEHeaderType), signature)) {
                InvalidPEFileException.ThrowInvalidPEHeader();
            }

            PEHeader hdr = new PEHeader();

            hdr.Type = signature;
            hdr.LinkerVersion = new Version(br.ReadByte(), br.ReadByte());
            hdr.SizeOfCode = br.ReadUInt32();
            hdr.SizeOfInitializedData = br.ReadUInt32();
            hdr.SizeOfUninitializedData = br.ReadUInt32();
            hdr.AddressOfEntryPoint = br.ReadUInt32();
            hdr.BaseOfCode = br.ReadUInt32();
            hdr.BaseOfData = br.ReadUInt32();

            if (signature == PEHeaderType.PE64) {
                hdr.ImageBase = br.ReadUInt64();
            }
            else {
                hdr.ImageBase = br.ReadUInt32();
            }

            hdr.SectionAlignment = br.ReadUInt32();
            hdr.FileAlignment = br.ReadUInt32();
            hdr.OperatingSystemVersion = new Version(br.ReadUInt16(), br.ReadUInt16());
            hdr.ImageVersion = new Version(br.ReadUInt16(), br.ReadUInt16());
            hdr.SubsystemVersion = new Version(br.ReadUInt16(), br.ReadUInt16());
            hdr.Win32VersionValue = br.ReadUInt32();
            hdr.SizeOfImage = br.ReadUInt32();
            hdr.SizeOfHeaders = br.ReadUInt32();
            hdr.Checksum = br.ReadUInt32();
            hdr.Subsystem = br.ReadUInt16();
            hdr.DllCharacteristics = br.ReadUInt16();

            if (signature == PEHeaderType.PE64) {
                hdr.SizeOfStackReserve = br.ReadUInt64();
                hdr.SizeOfStackCommit = br.ReadUInt64();
                hdr.SizeOfHeapReserve = br.ReadUInt64();
                hdr.SizeOfHeapCommit = br.ReadUInt64();
            }
            else {
                hdr.SizeOfStackReserve = br.ReadUInt32();
                hdr.SizeOfStackCommit = br.ReadUInt32();
                hdr.SizeOfHeapReserve = br.ReadUInt32();
                hdr.SizeOfHeapCommit = br.ReadUInt32();
            }

            hdr.LoaderFlags = br.ReadUInt32();

            hdr.DataDirectories = new PEDataDirectory[br.ReadUInt32()];
            for (int i = 0; i < hdr.DataDirectories.Length; i++) {
                hdr.DataDirectories[i] = br.ReadRecord<PEDataDirectory>();
            }

            return hdr;
        }
    }

    [BinaryRecord]
    public sealed class PEDataDirectory {

        public uint VirtualAddress;
        public uint Size;

        public PEDataDirectory() {
        }

        public override string ToString() {
            return string.Format("PEDataDirectory, RVA=0x{0:x}, Size=0x{1:x}",
                VirtualAddress, Size);
        }
    }

    public enum PEHeaderType : ushort {
        PE32 = 0x10b,
        PE64 = 0x20b,
    }
}
